package com.zf.nospring.rabbitmq.publish_confirm;

import com.rabbitmq.client.*;
import com.zf.nospring.rabbitmq.utils.RabbitUtils;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.HashMap;
import java.util.concurrent.*;

/**
 * 发布确认：
 * 开启发布确认后，发送消息后rabbitmq服务会返回一个发布的结果（true/false），我们称之为 发布确认，
 * 如果返回的确认是true就证明确认发送成功了，如果是false则确认发送失败，可以进行重试或者做其他操作。
 * <p>
 * 发布成功和失败的定义：
 * - 确认发布成功： "消息被投递到了匹配的队列上，如果消息是持久化消息，还要等消息写入磁盘" 才算是发布成功，对此rabbitmq就会返回一个确认成功（true）给生产者。
 * - 确认发布失败：消息没有被投递到匹配的队列上，或者出计划消息写入磁盘失败，那么就是发布失败，则rabbitmq会返回一个确认失败（false）给生产者。
 * <p>
 * rabbitmq有三种发布确认模式，分别为：
 * 1.单条确认模式：每发送一条消息，就获取一次确认结果，发布相同数量的消息，这种确认模式耗时最长；该方式同步阻塞发送过程。
 * 2.批量确认模式：每发送一批消息，就获取一次确认结果，若有消息无法发出，该模式无法确认是哪个消息无法发送；该方式也同步阻塞发送过程。
 * 3.异步确认模式：使用回调异步获取确认结果，该模式性能最好，在有错误情况下很好处理。性能最好。该方式不阻塞发送过程。
 * <p>
 * 注意：这三种模式都是要满足 "消息被投递到了匹配的队列上，如果消息是持久化消息，还要等消息写入磁盘" 才算是发布成功，否则发布失败。
 */
public class PublishConfirmProducer {
    public static final String testPublishConfirmQueue = "test_publish_confirm_queue";
    public static final HashMap<String, Object> queueArgs = new HashMap<>();

    static {
        // 配置如下参数来产生发送失败的消息
        // 指定队列的参数，指定队列最大只能存储 2 个消息
        queueArgs.put("x-max-length", 2);
        // "x-overflow"="reject-publish" 表示队列中的消息超过 "x-max-length" 个后，新来的消息将被丢弃，
        // 如果不配置，则默认是队列中的消息超过 "x-max-length" 个后删除队列的第一个消息，也就是最早到达队列的消息，然后写入新消息。
        queueArgs.put("x-overflow", "reject-publish");
    }


    public static void main(String[] args) throws IOException, TimeoutException, InterruptedException {
        // singleConfirm();
        // batchConfirm();
        asyncConfirm();
    }

    /**
     * 单个确认模式：
     * 每发送一条消息，就获取一次确认结果，发布相同数量的消息，这种确认模式耗时最长；该方式同步阻塞发送过程
     */
    public static void singleConfirm() throws IOException, TimeoutException, InterruptedException {
        Connection connection = RabbitUtils.getDefaultConnection();
        Channel channel = connection.createChannel();

        // 声明队列
        channel.queueDeclare(testPublishConfirmQueue, true, false, false, queueArgs);

        // 启用该通道的发布确认
        channel.confirmSelect();

        // 发送消息，每发送一次就等待一下确认结果
        for (int i = 1; i <= 5; i++) {
            // 发送消息
            String msg = "消息" + i;
            channel.basicPublish("", testPublishConfirmQueue, null, msg.getBytes(StandardCharsets.UTF_8));
            // 获取rabbitmq返回的确认结果，阻塞在此处，直到返回结果，成功 true，失败 false
            boolean res = channel.waitForConfirms();
            // 超时等待，获取rabbitmq返回的确认结果，阻塞在此处，直到返回结果或者超时抛出异常
            // boolean res = channel.waitForConfirms(2000);
            if (res) {
                System.out.printf("%s 发送成功\n", msg);
            } else {
                System.out.printf("%s 发送失败\n", msg);
            }
        }

        channel.close();
        connection.close();
    }

    /**
     * 批量确认模式：
     * 每发送一批消息，就获取一次确认结果，若有消息无法发出，该模式无法确认是哪个消息无法发送；该方式同步阻塞发送过程
     */
    public static void batchConfirm() throws IOException, TimeoutException, InterruptedException {
        Connection connection = RabbitUtils.getDefaultConnection();
        Channel channel = connection.createChannel();

        channel.queueDeclare(testPublishConfirmQueue, true, false, false, queueArgs);

        // 启用该通道的发布确认
        channel.confirmSelect();

        int count = 0, batchSize = 5;
        // 发送消息
        for (int i = 1; i <= 20; i++) {
            channel.basicPublish("", testPublishConfirmQueue, null, ("消息" + i).getBytes(StandardCharsets.UTF_8));
            count++;

            // 每发送5条消息就获取一次确认结果
            if (count == batchSize) {
                // 获取rabbitmq返回的确认结果，阻塞在此处，直到返回结果，成功 true，失败 false
                boolean res = channel.waitForConfirms();
                // 超时等待，获取rabbitmq返回的确认结果，阻塞在此处，直到返回结果或者超时抛出异常
                // boolean res = channel.waitForConfirms(2000);
                if (res) {
                    System.out.println("发送成功");
                } else {
                    System.out.println("发送失败，但是不知道那条消息失败了");
                }
                count = 0;
            }
        }

        // 如果还剩余一些消息，不足5条，再次对最后剩余的消息获取确认结果
        if (count > 0) {
            // 获取rabbitmq返回的确认结果，阻塞在此处，直到返回结果，成功 true，失败 false
            boolean res = channel.waitForConfirms();
            // 超时等待，获取rabbitmq返回的确认结果，阻塞在此处，直到返回结果或者超时抛出异常
            // boolean res = channel.waitForConfirms(2000);
            if (res) {
                System.out.println("发送成功");
            } else {
                System.out.println("发送失败，但是不知道那条消息失败了");
            }
            count = 0;
        }

        channel.close();
        connection.close();
    }


    /**
     * 异步确认：
     * 使用回调异步获取确认结果，该模式性能最好，在有错误情况下很好处理，性能最好。该方式不阻塞发送过程
     */
    public static void asyncConfirm() throws IOException, TimeoutException, InterruptedException {
        Connection connection = RabbitUtils.getDefaultConnection();
        Channel channel = connection.createChannel();
        // 声明一个队列
        channel.queueDeclare(testPublishConfirmQueue, true, false, false, queueArgs);

        // 启用该通道的发布确认
        channel.confirmSelect();

        // 用跳表存储消息，跳表中的数据按照key有序, deliveryTag -> msg
        ConcurrentSkipListMap<Long, String> allMessages = new ConcurrentSkipListMap<>();
        // 存储发送失败的消息
        ConcurrentLinkedQueue<String> failedMessages = new ConcurrentLinkedQueue<>();

        // 添加异步回调，代码是异步执行的，消息确认有可能是批量确认的，是否批量确认在于返回的multiple的参数，
        // 此参数如果true表示rabbitmq批量确认了 <= deliveryTag 的所有消息，如果为false的话表示单条确认。
        channel.addConfirmListener(new ConfirmListener() {
            /**
             * 消息发送成功的回调
             * @param deliveryTag 发送消息时的唯一标识，基于通道的，从 1 开始
             * @param multiple 是否是批量确认发送成功的
             */
            @Override
            public void handleAck(long deliveryTag, boolean multiple) throws IOException {
                System.out.printf("deliveryTag：%d 消息发送成功了，是否批量：%b\n", deliveryTag, multiple);
                // 批量确认发送成功
                if (multiple) {
                    // 移除 小于等于deliveryTag 的消息，因为这些消息发送成功了
                    // headMap 获取跳表中 key 小于deliveryTag 元素的一个视图，true 表示包含当前的 deliveryTag，也就是 小于等于deliveryTag
                    ConcurrentNavigableMap<Long, String> headMap = allMessages.headMap(deliveryTag, true);
                    // 移除
                    headMap.clear();
                } else {
                    // 单条确认发送成功，移除发送成功的消息
                    allMessages.remove(deliveryTag);
                }
            }

            /**
             * 消息发送失败的回调，对于发送失败的消息，可以将其写入一个共享的队列中让发布线程再次发布或者做其他处理
             * @param deliveryTag 发送消息时的唯一标识，基于通道的，从1开始
             * @param multiple 是否是批量确认发送失败的
             */
            @Override
            public void handleNack(long deliveryTag, boolean multiple) throws IOException {
                System.out.printf("deliveryTag：%d 消息发送失败了，是否批量：%b\n", deliveryTag, multiple);
                // 批量确认发送失败
                if (multiple) {
                    //获取 小于等于 deliveryTag 的消息，这些消息发送失败了
                    ConcurrentNavigableMap<Long, String> headMap = allMessages.headMap(deliveryTag, true);
                    // 收集失败的消息
                    failedMessages.addAll(headMap.values());
                    // 清除这些失败的消息
                    headMap.clear();
                } else {
                    // 单条确认发送失败
                    // 收集失败的消息
                    failedMessages.add(allMessages.get(deliveryTag));
                    allMessages.remove(deliveryTag);
                }
            }
        });

        // 发送消息
        for (int i = 1; i <= 20; i++) {
            String msg = "消息" + i;
            // 将消息的deliveryTag 及 消息 存入map中，channel.getNextPublishSeqNo() 获取即将发送的消息的 deliveryTag
            allMessages.put(channel.getNextPublishSeqNo(), msg);
            channel.basicPublish("", testPublishConfirmQueue, null, msg.getBytes(StandardCharsets.UTF_8));
        }

        // 异步回调原理：
        // 每个Connection(AMQConnection对象)连接中都会创建一个单独 MainLoop 线程，内部有循环不断读取一个帧，如果不是心跳帧，就发送到该帧所引用的通道，
        // 使通道执行组装好的命令，执行命令的过程中会执行异步操作，我们的 ConfirmListener 回调也是异步操作，就是在该线程内部执行的
        // 对于 ConfirmListener 回调：线程每次循环要么执行 ConfirmListener.handleAck(), 要么执行 ConfirmListener.handleNack()，

        // 所以不能立即关闭通道和连接，否则MainLoop线程就会终止，无法执行到回调。
        // 所以可以在此处调用 channel.waitForConfirms() 阻塞等待所有的消息都异步确认完毕，然后再关闭通道和连接，不用关心返回值，异步回调会处理成功或者失败。
        channel.waitForConfirms();

        // 所有的消息都确认完毕后，查看是否有消息发送失败了
        if (!failedMessages.isEmpty()) {
            System.out.println("有消息发送失败了，请处理: ");
            System.out.println(failedMessages);
        } else {
            System.out.println("所有的消息都发送成功了");
        }

        channel.close();
        connection.close();
    }
}
